
import os, glob, random, math
from PIL import Image
import numpy as np
import torch
from torch.utils.data import Dataset
from torchvision.transforms.functional import to_tensor
from utils.common import downsample_bicubic

class SRHRDataset(Dataset):
    def __init__(self, root, split="train", hr_crop=128, scale_min=1.1, scale_max=4.0):
        super().__init__()
        assert split in ["train", "val"]
        self.hr_dir = os.path.join(root, "HR", split)
        self.files = sorted(glob.glob(os.path.join(self.hr_dir, "*.png")) + 
                            glob.glob(os.path.join(self.hr_dir, "*.jpg")) + 
                            glob.glob(os.path.join(self.hr_dir, "*.jpeg")))
        if len(self.files) == 0:
            raise FileNotFoundError(f"No HR images found under {self.hr_dir}")
        self.hr_crop = hr_crop
        self.scale_min = scale_min
        self.scale_max = scale_max
        self.split = split

    def __len__(self):
        return len(self.files)

    def _random_hr_crop(self, img):
        W, H = img.size
        if H < self.hr_crop or W < self.hr_crop:
            pad_h = max(0, self.hr_crop - H)
            pad_w = max(0, self.hr_crop - W)
            img = np.array(img)
            if pad_h>0:
                img = np.pad(img, ((0,pad_h),(0,0),(0,0)), mode="reflect")
            if pad_w>0:
                img = np.pad(img, ((0,0),(0,pad_w),(0,0)), mode="reflect")
            img = Image.fromarray(img.astype(np.uint8))
            W, H = img.size
        x = random.randint(0, W - self.hr_crop)
        y = random.randint(0, H - self.hr_crop)
        return img.crop((x, y, x + self.hr_crop, y + self.hr_crop))

    def __getitem__(self, idx):
        path = self.files[idx]
        hr_img = Image.open(path).convert("RGB")
        if self.split == "train":
            hr_img = self._random_hr_crop(hr_img)
        hr = to_tensor(hr_img)
        Hh, Wh = hr.shape[1], hr.shape[2]

        s = random.uniform(self.scale_min, self.scale_max)
        lr = downsample_bicubic(hr, s)

        data = {
            "lr": lr,
            "hr": hr,
            "scale": s,
            "hr_size": (Hh, Wh),
            "path": path
        }
        return data
